import torch
import torch.nn as nn
import torch.optim as optim

idx2char = ['e', 'h', 'l', 'o']
char2idx = {ch: i for i, ch in enumerate(idx2char)}  # {'e': 0, 'h': 1, 'l': 2, 'o': 3}

x_data = [[1, 0, 2, 2]]
y_data = [0, 2, 2, 3]
# PyTorch规范要求：
# nn.Embedding层的输入必须是LongTensor类型，因为它是用来查找嵌入表中的索引
# 分类任务的标签(labels)通常也要求是长整型，特别是配合nn.CrossEntropyLoss使用时
inputs = torch.LongTensor(x_data)
labels = torch.LongTensor(y_data)


class CharRNN(nn.Module):
    def __init__(self, vocab_size, embed_dim=10, hidden_size=16, num_layers=1):
        super().__init__()
        self.emb = nn.Embedding(vocab_size, embed_dim)
        self.rnn = nn.RNN(
            input_size=embed_dim,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True
        )
        self.fc = nn.Linear(hidden_size, vocab_size)

    def forward(self, x):
        # 经过 x = self.emb(x) 的计算，x 的维度确实发生了变化：
        # 输入维度: (B, T)
        # 输出维度: (B, T, E)，其中 E 是嵌入维度(embed_dim)
        # B (Batch Size): 批次大小，表示一次处理多少个样本
        # T (Time Steps/Sequence Length): 时间步长/序列长度，表示每个样本包含多少个时间步
        # E (Embedding Dimension): 嵌入维度，表示每个token被映射到的向量维度
        x = self.emb(x)  # (B, T, E)
        # out: many to many
        # hn: many to one
        out, hn = self.rnn(x)  # (B, T, H)
        logits = self.fc(out)  # (B, T, V)
        return logits


vocab_size = len(idx2char)
model = CharRNN(vocab_size=vocab_size)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.05)

model.train()

epochs = 200
for epoch in range(1, epochs + 1):
    optimizer.zero_grad()
    logits = model(inputs)  # (B, T, V)
    # logits.reshape(-1, vocab_size): 将模型输出的 logits 调整为二维形状，第一维是所有位置的 token，第二维是词汇表大小
    # labels.reshape(-1): 将标签调整为一维形状，与 logits 的第一个维度对应
    # criterion(): 使用损失函数（通常是交叉熵）计算预测值与真实标签之间的损失
    # logits: 假设形状为 (B, T, V)，其中 B 是batch大小
    # logits.reshape(-1, vocab_size): 将其重塑为 (B*T, V)，即将所有时间步展平
    # labels: 假设形状为 (B, T)
    # labels.reshape(-1): 将其重塑为 (B*T,)，同样将所有时间步展平
    loss = criterion(logits.reshape(-1, vocab_size), labels.reshape(-1))
    # print(logits.shape, labels.shape)
    # print("logits.reshape(-1, vocab_size):", logits.reshape(-1, vocab_size))
    # print("labels.reshape(-1):", labels.reshape(-1))
    loss.backward()
    optimizer.step()
    if epoch % 50 == 0:
        with torch.no_grad():
            # 获取预测结果：从模型输出的 logits 中找出每个位置概率最大的类别索引
            # argmax(dim=-1)：在最后一个维度（词汇表维度）上找到最大值的索引
            # 输出形状：(B, T)，即批次大小 × 序列长度的预测索引矩阵
            print("logits.shape:", logits.shape)
            print("logits:", logits)
            pred = logits.argmax(dim=-1)  # (B, T)
            print("pred.shape:", pred.shape)
            print("pred:", pred)
            pred_str = ''.join(idx2char[i] for i in pred[0].tolist())
        print(f"[epoch {epoch:03d}] loss={loss.item():.4f} | pred={pred_str}")

save_path = "../models/char_rnn_state_dict.pth"
torch.save(model.state_dict(), save_path)
print(f"模型参数已保存到: {save_path}")
